module test_cuda_02_gpu_const

  implicit none

  integer, parameter :: r8 = 8
  integer, parameter :: ns = 1
  integer, parameter :: nblk = 80
  integer, parameter :: nthr = 512
  integer, parameter :: nx = nblk * nthr
  integer, parameter :: nt = 1000000
  real(r8), parameter :: dx = 1.0d0 / nx
  real(r8), parameter :: c = 1.0d0
  real(r8), parameter :: dt = 0.8d0 * dx / c

end module test_cuda_02_gpu_const

module test_cuda_02_gpu_kernels

  use cudafor
  use test_cuda_02_gpu_const

contains

  attributes(global) subroutine upwind(q, f)

    real(r8), device :: q(1-ns:nx+ns)
    real(r8), device :: f(1-ns:nx+ns)

    integer i

    i = (blockIdx%x - 1) * blockDim%x + threadIdx%x
    f(i) = 0.5d0 * (c * (q(i+1) + q(i)) - abs(c) * (q(i+1) - q(i)))

  end subroutine upwind

  attributes(global) subroutine update(f, q_old, q_new)

    real(r8), device :: f(1-ns:nx+ns)
    real(r8), device :: q_old(1-ns:nx+ns)
    real(r8), device :: q_new(1-ns:nx+ns)

    integer i

    i = (blockIdx%x - 1) * blockDim%x + threadIdx%x
    q_new(i) = q_old(i) - dt / dx * (f(i) - f(i-1))

  end subroutine update

  attributes(global) subroutine fill_halo(x)

    real(r8), device :: x(1-ns:nx+ns)

    integer i

    i = (blockIdx%x - 1) * blockDim%x + threadIdx%x
    if (i <= ns) then
      ! Left halo
      x(1-i) = x(nx+1-i)
      ! Right halo
      x(nx+i) = x(i)
    end if

  end subroutine fill_halo

end module test_cuda_02_gpu_kernels

program test_cuda_02_gpu

  use cudafor
  use test_cuda_02_gpu_const
  use test_cuda_02_gpu_kernels

  implicit none

  real(r8), allocatable :: x_h(:)
  real(r8), allocatable :: q_h(:)
  real(r8), allocatable, device :: q_d(:,:)
  real(r8), allocatable, device :: f_d(:)

  integer i, old, new, ierr

  allocate(x_h(1-ns:nx+ns))
  allocate(q_h(1-ns:nx+ns))
  allocate(q_d(1-ns:nx+ns,2))
  allocate(f_d(1-ns:nx+ns))

  do i = lbound(x_h, 1), ubound(x_h, 1)
    x_h(i) = (i - 1) * dx
    if (x_h(i) >= 0.05 .and. x_h(i) <= 0.3) then
      q_h(i) = 1
    else
      q_h(i) = 0
    end if
  end do

  q_d(:,1) = q_h

  old = 1; new = 2
  do i = 1, nt
    call upwind<<<nblk,nthr>>>(q_d(:,old), f_d)
    call fill_halo<<<1,ns>>>(f_d)
    ierr = cudaDeviceSynchronize()
    call update<<<nblk,nthr>>>(f_d, q_d(:,old), q_d(:,new))
    call fill_halo<<<1,ns>>>(q_d(:,new))
    ierr = cudaDeviceSynchronize()
    call time_advance()
  end do

  q_h = q_d(:,old)

  print *, sum(x_h(1:nx) * q_h(1:nx))

  deallocate(x_h)
  deallocate(q_h)
  deallocate(q_d)
  deallocate(f_d)

contains

  subroutine time_advance()

    integer tmp

    tmp = old
    old = new
    new = tmp

  end subroutine time_advance

end program test_cuda_02_gpu
